{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "介绍如何在pytorch环境下，使用CW2算法攻击基于MNIST数据集预训练的CNN/MLP模型。运行该文件前，需要先运行指定文件生成对应的模型：\n",
    "\n",
    "    cd tutorials\n",
    "    python mnist_model_pytorch.py\n",
    "     "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Jupyter notebook中使用Anaconda中的环境需要单独配置，默认情况下使用的是系统默认的Python环境，以使用advbox环境为例。\n",
    "首先在默认系统环境下执行以下命令，安装ipykernel。\n",
    "\n",
    "    conda install ipykernel\n",
    "    conda install -n advbox ipykernel\n",
    "\n",
    "在advbox环境下激活，这样启动后就可以在界面上看到advbox了。\n",
    "\n",
    "    python -m ipykernel install --user --name advbox --display-name advbox \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_data: torch.Size([60000, 28, 28])\n",
      "train_labels: torch.Size([60000])\n",
      "test_data: torch.Size([10000, 28, 28])\n"
     ]
    }
   ],
   "source": [
    "#调试开关\n",
    "import logging\n",
    "#logging.basicConfig(level=logging.INFO,format=\"%(filename)s[line:%(lineno)d] %(levelname)s %(message)s\")\n",
    "#logger=logging.getLogger(__name__)\n",
    "import sys\n",
    "import torch\n",
    "import torchvision\n",
    "from torchvision import datasets, transforms\n",
    "from torch.autograd import Variable\n",
    "import torch.utils.data.dataloader as Data\n",
    "from adversarialbox.adversary import Adversary\n",
    "from adversarialbox.attacks.cw2_pytorch import CW_L2_Attack\n",
    "from adversarialbox.models.pytorch import PytorchModel\n",
    "from tutorials.mnist_model_pytorch import Net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n",
      "attack success, original_label=9, adversarial_label=4, count=1\n",
      "attack success, original_label=3, adversarial_label=2, count=2\n",
      "attack success, original_label=1, adversarial_label=7, count=3\n",
      "attack success, original_label=5, adversarial_label=3, count=4\n",
      "attack success, original_label=0, adversarial_label=2, count=5\n",
      "attack success, original_label=6, adversarial_label=4, count=6\n",
      "attack success, original_label=2, adversarial_label=8, count=7\n",
      "attack success, original_label=1, adversarial_label=4, count=8\n",
      "attack success, original_label=8, adversarial_label=9, count=9\n",
      "attack success, original_label=1, adversarial_label=6, count=10\n",
      "attack success, original_label=2, adversarial_label=1, count=11\n",
      "attack success, original_label=4, adversarial_label=9, count=12\n",
      "attack success, original_label=8, adversarial_label=3, count=13\n",
      "attack success, original_label=7, adversarial_label=9, count=14\n",
      "attack success, original_label=6, adversarial_label=5, count=15\n",
      "attack success, original_label=1, adversarial_label=4, count=16\n",
      "attack success, original_label=5, adversarial_label=9, count=17\n",
      "attack success, original_label=7, adversarial_label=2, count=18\n",
      "attack success, original_label=8, adversarial_label=2, count=19\n",
      "attack success, original_label=1, adversarial_label=4, count=20\n",
      "attack success, original_label=1, adversarial_label=4, count=21\n",
      "attack success, original_label=4, adversarial_label=7, count=22\n",
      "attack success, original_label=7, adversarial_label=3, count=23\n",
      "attack success, original_label=9, adversarial_label=4, count=24\n",
      "attack success, original_label=7, adversarial_label=9, count=25\n",
      "attack success, original_label=6, adversarial_label=5, count=26\n",
      "attack success, original_label=3, adversarial_label=9, count=27\n",
      "attack success, original_label=1, adversarial_label=4, count=28\n",
      "attack success, original_label=2, adversarial_label=7, count=29\n",
      "attack success, original_label=1, adversarial_label=7, count=30\n",
      "attack success, original_label=0, adversarial_label=9, count=31\n",
      "attack success, original_label=7, adversarial_label=3, count=32\n",
      "attack success, original_label=7, adversarial_label=3, count=33\n",
      "attack success, original_label=8, adversarial_label=2, count=34\n",
      "attack success, original_label=5, adversarial_label=3, count=35\n",
      "attack success, original_label=0, adversarial_label=6, count=36\n",
      "attack success, original_label=9, adversarial_label=4, count=37\n",
      "attack success, original_label=4, adversarial_label=9, count=38\n",
      "attack success, original_label=1, adversarial_label=4, count=39\n",
      "attack success, original_label=5, adversarial_label=3, count=40\n",
      "attack success, original_label=6, adversarial_label=5, count=41\n",
      "attack success, original_label=0, adversarial_label=6, count=42\n",
      "attack success, original_label=4, adversarial_label=9, count=43\n",
      "attack success, original_label=6, adversarial_label=5, count=44\n",
      "attack success, original_label=6, adversarial_label=0, count=45\n",
      "attack success, original_label=4, adversarial_label=9, count=46\n",
      "attack success, original_label=9, adversarial_label=4, count=47\n",
      "attack success, original_label=3, adversarial_label=7, count=48\n",
      "attack success, original_label=5, adversarial_label=3, count=49\n",
      "attack success, original_label=8, adversarial_label=9, count=50\n",
      "attack success, original_label=4, adversarial_label=9, count=51\n",
      "attack success, original_label=2, adversarial_label=7, count=52\n",
      "attack success, original_label=8, adversarial_label=2, count=53\n",
      "attack success, original_label=2, adversarial_label=3, count=54\n",
      "attack success, original_label=0, adversarial_label=6, count=55\n",
      "attack success, original_label=4, adversarial_label=2, count=56\n",
      "attack success, original_label=5, adversarial_label=8, count=57\n",
      "attack success, original_label=3, adversarial_label=5, count=58\n",
      "attack success, original_label=3, adversarial_label=5, count=59\n",
      "attack success, original_label=3, adversarial_label=8, count=60\n",
      "attack success, original_label=1, adversarial_label=4, count=61\n",
      "attack success, original_label=1, adversarial_label=4, count=62\n",
      "attack success, original_label=9, adversarial_label=4, count=63\n",
      "attack success, original_label=4, adversarial_label=8, count=64\n",
      "attack success, original_label=2, adversarial_label=3, count=65\n",
      "attack success, original_label=6, adversarial_label=0, count=66\n",
      "attack success, original_label=1, adversarial_label=6, count=67\n",
      "attack success, original_label=5, adversarial_label=3, count=68\n",
      "attack success, original_label=4, adversarial_label=6, count=69\n",
      "attack success, original_label=2, adversarial_label=3, count=70\n",
      "attack success, original_label=0, adversarial_label=2, count=71\n",
      "attack success, original_label=1, adversarial_label=7, count=72\n",
      "attack success, original_label=7, adversarial_label=9, count=73\n",
      "attack success, original_label=7, adversarial_label=9, count=74\n",
      "attack success, original_label=8, adversarial_label=2, count=75\n",
      "attack success, original_label=1, adversarial_label=4, count=76\n",
      "attack success, original_label=4, adversarial_label=8, count=77\n",
      "attack success, original_label=1, adversarial_label=4, count=78\n",
      "attack success, original_label=2, adversarial_label=7, count=79\n",
      "attack success, original_label=6, adversarial_label=2, count=80\n",
      "attack success, original_label=9, adversarial_label=4, count=81\n",
      "attack success, original_label=0, adversarial_label=2, count=82\n",
      "attack success, original_label=2, adversarial_label=8, count=83\n",
      "attack success, original_label=5, adversarial_label=3, count=84\n",
      "attack success, original_label=8, adversarial_label=5, count=85\n",
      "attack success, original_label=1, adversarial_label=4, count=86\n",
      "attack success, original_label=4, adversarial_label=9, count=87\n",
      "attack success, original_label=0, adversarial_label=5, count=88\n",
      "attack success, original_label=1, adversarial_label=6, count=89\n",
      "attack success, original_label=0, adversarial_label=8, count=90\n",
      "attack success, original_label=2, adversarial_label=3, count=91\n",
      "attack success, original_label=4, adversarial_label=8, count=92\n",
      "attack success, original_label=8, adversarial_label=3, count=93\n",
      "attack success, original_label=1, adversarial_label=4, count=94\n",
      "attack success, original_label=0, adversarial_label=9, count=95\n",
      "attack success, original_label=1, adversarial_label=4, count=96\n",
      "attack success, original_label=5, adversarial_label=3, count=97\n",
      "attack success, original_label=7, adversarial_label=1, count=98\n",
      "attack success, original_label=1, adversarial_label=7, count=99\n",
      "attack success, original_label=9, adversarial_label=4, count=100\n",
      "[TEST_DATASET]: fooling_count=100, total_count=100, fooling_rate=1.000000\n",
      "cw2 attack done\n"
     ]
    }
   ],
   "source": [
    "TOTAL_NUM = 100\n",
    "pretrained_model=\"tutorials/mnist-pytorch/net.pth\"\n",
    "loss_func = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "#使用MNIST测试数据集 随机挑选TOTAL_NUM个\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST('tutorials/mnist-pytorch/data', train=False, download=True, transform=transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "    ])),\n",
    "    batch_size=1, shuffle=True)\n",
    "\n",
    "# Define what device we are using\n",
    "logging.info(\"CUDA Available: {}\".format(torch.cuda.is_available()))\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "#这里有个需要注意的地方 cw的输出必须是logit层而不能是softmax层 否则极大概率梯度消失一直无法收敛\n",
    "# Initialize the network\n",
    "model = Net().to(device)\n",
    "\n",
    "# Load the pretrained model\n",
    "model.load_state_dict(torch.load(pretrained_model, map_location='cpu'))\n",
    "\n",
    "# Set the model in evaluation mode. In this case this is for the Dropout layers\n",
    "model.eval()\n",
    "\n",
    "# advbox demo\n",
    "m = PytorchModel(\n",
    "    model, loss_func,(0.0, 1.0),\n",
    "    channel_axis=1)\n",
    "\n",
    "#实例化CW_L2_Attack\n",
    "attack = CW_L2_Attack(m)\n",
    "#设置分类数num_labels 最大迭代次数max_iterations 二分查找次数 binary_search_steps C初始化值 initial_const\n",
    "attack_config = {\"num_labels\": 10,\"max_iterations\":1000,\"binary_search_steps\":4,\"initial_const\":100.0}\n",
    "\n",
    "# use test data to generate adversarial examples\n",
    "total_count = 0\n",
    "fooling_count = 0\n",
    "\n",
    "for i, data in enumerate(test_loader):\n",
    "    inputs, labels = data\n",
    "    inputs, labels=inputs.numpy(),labels.numpy()\n",
    "\n",
    "    total_count += 1\n",
    "    adversary = Adversary(inputs, labels[0])\n",
    "\n",
    "    # FGSM non-targeted attack\n",
    "    adversary = attack(adversary, **attack_config)\n",
    "\n",
    "    if adversary.is_successful():\n",
    "        fooling_count += 1\n",
    "        print(\n",
    "            'attack success, original_label=%d, adversarial_label=%d, count=%d'\n",
    "            % (labels, adversary.adversarial_label, total_count))\n",
    "\n",
    "    else:\n",
    "        print('attack failed, original_label=%d, count=%d' %\n",
    "              (labels, total_count))\n",
    "\n",
    "    if total_count >= TOTAL_NUM:\n",
    "        print(\n",
    "            \"[TEST_DATASET]: fooling_count=%d, total_count=%d, fooling_rate=%f\"\n",
    "            % (fooling_count, total_count,\n",
    "               float(fooling_count) / total_count))\n",
    "        break\n",
    "print(\"cw2 attack done\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "advbox",
   "language": "python",
   "name": "advbox"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
